package cryption

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"errors"
	"fmt"
	"strings"
)

const RSA_ALGORITHM_KEY_TYPE = "PKCS8"

/**
 * @description: RSA公钥私钥产生
 * @return {string, string, error} RSA私钥, RSA公钥
 */
func GenRsaKey() (prvkey, pubkey string, err error) {
	// 生成私钥文件
	privateKey, err := rsa.GenerateKey(rand.Reader, 1024)
	if err != nil {
		return "", "", err
	}

	// 根据类型生成私钥
	var derStream []byte
	switch RSA_ALGORITHM_KEY_TYPE {
	case "PKCS1":
		derStream = x509.MarshalPKCS1PrivateKey(privateKey)
	case "PKCS8":
		derStream, err = x509.MarshalPKCS8PrivateKey(privateKey)
		if err != nil {
			return "", "", err
		}
	}
	block := &pem.Block{
		Type:  "PRIVATE KEY",
		Bytes: derStream,
	}
	prvkey = string(pem.EncodeToMemory(block))
	publicKey := &privateKey.PublicKey
	derPkix, err := x509.MarshalPKIXPublicKey(publicKey)
	if err != nil {
		return "", "", err
	}
	block = &pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: derPkix,
	}
	pubkey = string(pem.EncodeToMemory(block))

	return prvkey, pubkey, nil
}

/**
 * @description: 签名
 * @param {string} text 需要签名的数据
 * @param {string} prvkey 签名RSA私钥
 * @return {string, error} base64编码后的字符串
 */
func RsaSignWithSha256(text, prvkey string) (string, error) {
	h := sha256.New()
	data := []byte(text)
	h.Write(data)
	hashed := h.Sum(nil)
	keyBytes := []byte(prvkey)
	block, _ := pem.Decode(keyBytes)
	if block == nil {
		return "", errors.New("private key error")
	}

	// 根据类型解析私钥
	var priv *rsa.PrivateKey
	switch RSA_ALGORITHM_KEY_TYPE {
	case "PKCS1":
		var err error
		priv, err = x509.ParsePKCS1PrivateKey(block.Bytes)
		if err != nil {
			fmt.Println("ParsePKCS8PrivateKey err", err)
			return "", err
		}
	case "PKCS8":
		privateKey, err := x509.ParsePKCS8PrivateKey(block.Bytes)
		if err != nil {
			return "", err
		}
		priv, _ = privateKey.(*rsa.PrivateKey)
	}

	signature, err := rsa.SignPKCS1v15(rand.Reader, priv, crypto.SHA256, hashed)
	if err != nil {
		fmt.Printf("Error from signing: %s\n", err)
		return "", err
	}

	return base64.StdEncoding.EncodeToString(signature), nil
}

/**
 * @description: 验证签名
 * @param {string} text 验证签名的数据
 * @param {string} sign 签名base64编码后的字符串
 * @param {string} pubkey RSA私钥
 * @return {bool, error}
 */
func RsaVerySignWithSha256(text, sign, pubkey string) (bool, error) {
	keyBytes := []byte(pubkey)
	block, _ := pem.Decode(keyBytes)
	if block == nil {
		return false, errors.New("public key error")
	}
	pubKey, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return false, err
	}

	data := []byte(text)
	hashed := sha256.Sum256(data)
	signData, err := base64.StdEncoding.DecodeString(sign)
	if err != nil {
		return false, err
	}

	err = rsa.VerifyPKCS1v15(pubKey.(*rsa.PublicKey), crypto.SHA256, hashed[:], signData)
	if err != nil {
		return false, err
	}
	return true, nil
}

/**
 * @description: 公钥加密
 * @param {string} plaintext 加密数据
 * @param {string} pubKey RSA公钥
 * @return {string, error} 加密后字符串
 */
func RsaEncrypt(plaintext, pubKey string) (string, error) {
	// 解密pem格式的公钥
	keyBytes := []byte(pubKey)
	block, _ := pem.Decode(keyBytes)
	if block == nil {
		return "", errors.New("public key error")
	}
	// 解析公钥
	pubInterface, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return "", err
	}
	// 类型断言
	pub := pubInterface.(*rsa.PublicKey)
	// 加密
	data := []byte(plaintext)
	ciphertext, err := rsa.EncryptPKCS1v15(rand.Reader, pub, data)
	if err != nil {
		return "", err
	}

	return base64.StdEncoding.EncodeToString(ciphertext), nil
}

/**
 * @description: 私钥解密
 * @param {*} ciphertext 解密数据
 * @param {string} prvkey RSA密钥
 * @return {string, error} 解密后字符串
 */
func RsaDecrypt(ciphertext, prvkey string) (string, error) {
	// 获取私钥
	keyBytes := []byte(prvkey)
	block, _ := pem.Decode(keyBytes)
	if block == nil {
		return "", errors.New("private key error")
	}

	// 根据类型解析私钥
	var priv *rsa.PrivateKey
	switch RSA_ALGORITHM_KEY_TYPE {
	case "PKCS1":
		var err error
		priv, err = x509.ParsePKCS1PrivateKey(block.Bytes)
		if err != nil {
			fmt.Println("ParsePKCS8PrivateKey err", err)
			return "", err
		}
	case "PKCS8":
		privateKey, err := x509.ParsePKCS8PrivateKey(block.Bytes)
		if err != nil {
			return "", err
		}
		priv, _ = privateKey.(*rsa.PrivateKey)
	}

	// 解密
	enData, err := base64.StdEncoding.DecodeString(ciphertext)
	if err != nil {
		return "", err
	}

	data, err := rsa.DecryptPKCS1v15(rand.Reader, priv, enData)
	if err != nil {
		return "", err
	}

	return string(data), nil
}

/**
 * @description: 按长度插入字符串
 * @param {string} s 需要处理的字符串
 * @param {int} chunkSize 宽度
 * @param {string} sep 插入的字符
 * @return {string}
 */
func ChunkString(s string, chunkSize int, sep string) string {
	var chunks []string
	runes := []rune(s)

	if len(runes) == 0 {
		return s
	}

	for i := 0; i < len(runes); i += chunkSize {
		nn := i + chunkSize
		if nn > len(runes) {
			nn = len(runes)
		}
		chunks = append(chunks, string(runes[i:nn]))
	}

	return strings.Join(chunks, sep)
}

/**
 * @description: 初始化公钥
 * @param {string} pubKey 待初始化公钥
 * @param {bool} clear 是否去掉头尾
 * @return {string}
 */
func InitPublicKey(pubKey string, clear bool) string {
	pubKey = strings.ReplaceAll(pubKey, "-----BEGIN PUBLIC KEY-----", "")
	pubKey = strings.ReplaceAll(pubKey, "-----END PUBLIC KEY-----", "")
	pubKey = strings.ReplaceAll(pubKey, `\n`, "")
	pubKey = strings.ReplaceAll(pubKey, "\n", "")

	if clear {
		return pubKey
	}

	pubKey = fmt.Sprintf("-----BEGIN PUBLIC KEY-----\n%s\n-----END PUBLIC KEY-----", ChunkString(pubKey, 64, "\n"))
	return pubKey
}
